#!/bin/bash

# 快速安装缺失依赖的脚本
# 用于在 qwen32b 环境中安装缺失的包

echo "========================================="
echo "安装缺失的依赖包"
echo "========================================="

# 检查是否在 conda 环境中
if [ -z "$CONDA_DEFAULT_ENV" ]; then
    echo "⚠ 警告: 未检测到 Conda 环境"
    echo "正在尝试激活 qwen32b 环境..."
    
    # 尝试初始化 conda
    if [ -f "$HOME/miniconda/etc/profile.d/conda.sh" ]; then
        source "$HOME/miniconda/etc/profile.d/conda.sh"
    elif [ -f "$HOME/anaconda3/etc/profile.d/conda.sh" ]; then
        source "$HOME/anaconda3/etc/profile.d/conda.sh"
    else
        echo "✗ 无法找到 conda，请先运行: source ~/.bashrc"
        exit 1
    fi
    
    # 尝试激活环境
    if conda activate qwen32b 2>/dev/null; then
        echo "✓ 已激活 qwen32b 环境"
    else
        echo "✗ 无法激活 qwen32b 环境"
        echo "请手动运行: conda activate qwen32b"
        exit 1
    fi
else
    echo "当前 Conda 环境: $CONDA_DEFAULT_ENV"
    if [ "$CONDA_DEFAULT_ENV" != "qwen32b" ]; then
        echo "⚠ 警告: 建议在 qwen32b 环境中运行"
        read -p "是否切换到 qwen32b 环境? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            conda activate qwen32b
        fi
    fi
fi

echo ""
echo "开始检查并安装依赖..."

# 检查 PyTorch
echo ""
echo "【1/3】检查 PyTorch..."
if python -c "import torch" 2>/dev/null; then
    TORCH_VERSION=$(python -c "import torch; print(torch.__version__)" 2>/dev/null)
    echo "✓ PyTorch 已安装: $TORCH_VERSION"
    
    # 检查 CUDA 支持
    CUDA_AVAILABLE=$(python -c "import torch; print(torch.cuda.is_available())" 2>/dev/null)
    if [ "$CUDA_AVAILABLE" = "True" ]; then
        echo "✓ CUDA 可用"
    else
        echo "⚠ CUDA 不可用，但 PyTorch 已安装"
        echo "  如果需要 GPU 支持，可能需要重新安装 PyTorch CUDA 版本"
    fi
else
    echo "✗ PyTorch 未安装"
    echo "正在安装 PyTorch (CUDA 12.1)..."
    pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
    if [ $? -eq 0 ]; then
        echo "✓ PyTorch 安装成功"
    else
        echo "✗ PyTorch 安装失败"
        exit 1
    fi
fi

# 检查 vLLM
echo ""
echo "【2/3】检查 vLLM..."
if python -c "import vllm" 2>/dev/null; then
    VLLM_VERSION=$(python -c "import vllm; print(vllm.__version__)" 2>/dev/null)
    echo "✓ vLLM 已安装: $VLLM_VERSION"
else
    echo "✗ vLLM 未安装"
    echo "正在安装 vLLM..."
    pip install vllm
    if [ $? -eq 0 ]; then
        echo "✓ vLLM 安装成功"
    else
        echo "✗ vLLM 安装失败"
        echo "  注意: vLLM 安装可能需要较长时间，请耐心等待"
        exit 1
    fi
fi

# 检查 ModelScope
echo ""
echo "【3/3】检查 ModelScope..."
if python -c "import modelscope" 2>/dev/null; then
    echo "✓ ModelScope 已安装"
else
    echo "⚠ ModelScope 未安装（可选，用于下载模型）"
    read -p "是否安装 ModelScope? (y/N): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        pip install modelscope
        if [ $? -eq 0 ]; then
            echo "✓ ModelScope 安装成功"
        else
            echo "✗ ModelScope 安装失败"
        fi
    fi
fi

echo ""
echo "========================================="
echo "安装完成！"
echo "========================================="
echo ""
echo "现在可以运行测试脚本："
echo "  python quick_test_qwen.py"
echo "  或"
echo "  bash test_qwen32b.sh"

